import time
import torch
from modelscope import snapshot_download
from ipex_llm import optimize_model
from qwen_vl_utils import process_vision_info
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor

class Qwen2VLChatModel:
    def __init__(self, model_name='Qwen/Qwen2-VL-2B-Instruct',
                 local_dir='C:\models\Qwen\Qwen2-VL-2B-Instruct',
                 min_pixels=256 * 28 * 28,
                 max_pixels=1280 * 28 * 28,
                 max_new_tokens=1024):
        self.model_dir = snapshot_download(model_name, local_dir=local_dir)
        self.model = Qwen2VLForConditionalGeneration.from_pretrained(self.model_dir,
                                                                     trust_remote_code=True,
                                                                     torch_dtype='auto',
                                                                     low_cpu_mem_usage=True,
                                                                     use_cache=True)
        self.model = optimize_model(self.model, low_bit='sym_int4', modules_to_not_convert=["visual"])
        self.model = self.model.half().to("xpu")
        self.processor = AutoProcessor.from_pretrained(self.model_dir, min_pixels=min_pixels, max_pixels=max_pixels)
        self.max_new_tokens = max_new_tokens

    def chat(self, image_path, prompt):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": image_path,
                    },
                    {"type": "text", "text": prompt},
                ],
            }
        ]
        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to('xpu')

        with torch.inference_mode():
            st = time.time()
            generated_ids = self.model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens
            )
            torch.xpu.synchronize()
            end = time.time()
            generated_ids = generated_ids.cpu()
            generated_ids = [
                output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
            ]

            response = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
            print(f'Inference time: {end - st} s')
            return response


if __name__ == '__main__':
    model = Qwen2VLChatModel(local_dir="C:\models\Qwen\Qwen2-VL-2B-Instruct")
    image_file = "../examples/rag_data/images/0d03955a-7846-4da5-abf4-e021ef856544.jpeg"
    # image_file = Path("C:/WorkSpace/jingxiang.ai/aipcagent/examples/rag_data/images/0d03955a-7846-4da5-abf4-e021ef856544.jpeg")
    result = model.chat(image_file, '描述一下这张图')
    print(result)
